# -*- coding: utf-8 -*-


"""
法狗狗selenium自动化爬虫
使用Selenium驱动浏览器进行登录并爬取数据
"""


import json
import re
import sys
import time

from bs4 import BeautifulSoup
from selenium import webdriver
from selenium.common.exceptions import TimeoutException
from selenium.webdriver.chrome.options import Options
from selenium.webdriver.common.by import By
from selenium.webdriver.support import expected_conditions as EC
from selenium.webdriver.support.ui import WebDriverWait


class FggSpider():
    '''法狗狗selenium爬虫'''

    def __init__(
        self, headless=False, no_image=False, ignore_certificate=False):
        '''初始化浏览器
        :param headless: 是否显示浏览器界面
        :param no_image: 是否加载页面图片
        :param ignore_certificate: 是否忽略证书
        '''
        chrome_options = Options()
        chrome_options.add_argument('--window-size=1400x900')
        if headless:
            chrome_options.set_headless()
        if ignore_certificate:
            chrome_options.add_argument('--ignore-certificate-errors')
        if no_image:
            prefs = {'profile.managed_default_content_settings.images': 2}
            chrome_options.add_experimental_option("prefs", prefs)
        self.browser = webdriver.Chrome(
            'D:/files/python/chromedriver.exe',
            chrome_options=chrome_options)            
        self.wait = WebDriverWait(self.browser, 10)
        self.EC = EC
    
    def search_init(self, phone_number='18856899065', cookies=None):
        '''初始化搜索 进入问答页面'''
        self.browser.get('https://www.fagougou.com/pc/?mkt=eldmh0340fd45')
        self.wait.until(EC.presence_of_element_located(
            (By.CSS_SELECTOR, '#content .J-bd > ul')
        ))

        # add cookies
        if cookies:
            if isinstance(cookies, dict):
                self.browser.add_cookie(cookies)
            if isinstance(cookies, list):
                for i in cookies:
                    self.browser.add_cookie(i)
            else:
                pass
            self.browser.refresh()
            self.wait.until(EC.presence_of_element_located(
                (By.CSS_SELECTOR, '#content .J-bd > ul')
            ))
         
        target = self.browser.find_element_by_css_selector(
            '#content .J-bd > ul:nth-of-type(1) > li:nth-of-type(1) > a')
        target.click()       

        # login
        try:
            self.wait.until(EC)
            self.browser.find_element_by_id('login')
        except:
            print('已登录')
        else:
            self.wait.until(EC.presence_of_element_located(
                (By.ID, 'popup-phone')))
            phone = self.browser.find_element_by_css_selector(
                '#login input.phone')
            phone.clear()
            phone.send_keys(phone_number)
    
            get_phone_code = self.browser.find_element_by_css_selector('#login #getCode')
            get_phone_code.click()
    
            phone_code = self.browser.find_element_by_css_selector('#login #code')
            phone_code.clear()
            code = input('输入验证码: ')
            phone_code.send_keys(code)
        
            submit = self.browser.find_element_by_css_selector('input.btn')
            submit.click()
    
        try:
            self.wait.until(self.EC.presence_of_element_located(
                (By.CSS_SELECTOR, '#msg-textarea')
            ))
        except TimeoutException:
            print('error')
            return
        else:
            print('已进入问答页面')
            self.browser.refresh()
            with open('cookies.json', 'w', encoding='utf8') as f:
                json.dump({'cookies': self.browser.get_cookies()}, f, ensure_ascii=False)
            return 1

    def send_question(self, question):
        '''发送问题'''
        question_box = self.browser.find_element_by_css_selector('#msg-textarea')
        question_box.clear()
        question_box.send_keys(question)

        question_btn = self.browser.find_element_by_css_selector('#send-msg')
        question_btn.click()
    
        try:
            self.wait.until_not(self.EC.presence_of_element_located(
                (By.CSS_SELECTOR,
                '#chat-content > div > div:last-of-type[class="message me clearfix"]')
            ))
        except TimeoutException:
            return
        else:
            has_answer = self.browser.find_element_by_css_selector(
                '#chat-content > div > div:last-of-type')
            if has_answer.get_attribute('class') == 'message you law mt02 clearfix':
                print('无答案返回，只有相关问题')
                return 1
            elif has_answer.get_attribute('class') == 'message you law clearfix':
                has_answer = has_answer.find_element_by_css_selector('div > div.feed-back')
                if has_answer.get_attribute('data-queryrecorditemid'):
                    return 2
                else:
                    return
            else:
                return
    
    def get_answer(self, status_code):
        result = {}
        if not status_code:
            return
        elif status_code == 2: # 有回答有相关推荐
            answer = self.browser.find_elements_by_css_selector(
                '#chat-content > div > div:last-of-type > div.content')
            result['answer'] = answer[0].text
        
            # 相关法条
            try:
                get_related_law = self.browser.find_element_by_css_selector(
                    '#chat-content > div > div:last-of-type .content > div.more i.icon-down')
            except:
                result['related_law'] = []
            else:
                get_related_law.click()
                related_law = self.browser.find_elements_by_css_selector(
                    '#chat-content > div > div:last-of-type .content > div.more dl')
                result['related_law'] = [i.text for i in related_law]
        
            # 相关问题
            try:
                get_related_answer = answer[1].find_element_by_css_selector('a')
            except:
                result['related_answer'] = []
            else:
                try:
                    get_related_answer.click()
                    self.wait.until_not(EC.presence_of_element_located(
                        (By.CSS_SELECTOR,
                        '#chat-content > div > div:last-of-type[class="message you law clearfix"]')
                    ))
                except:
                    result['related_answer'] = []
                else:
                    related_answer = self.browser.find_elements_by_css_selector(
                        '#chat-content > div > div:last-of-type li')
                    result['related_answer'] = [i.text for i in related_answer]
        else: # 无回答只有相关问题
            result['answer'], result['related_law'] = 'none', []
            try:
                related_answer = self.browser.find_elements_by_css_selector(
                    '#chat-content > div > div:last-of-type > div > ul > li')
            except:
                result['related_answer'] = []
            else:
                result['related_answer'] = [i.text for i in related_answer]
        return result
    
    def close(self):
        try:
            self.browser.close()
        except:
            pass


class QAFiles(object):

    def __init__(self, question_file, target_file=None):
        self.__fr = open(question_file, encoding='utf-8')
        self.__fw = open(target_file, 'w', encoding='utf-8')
    
    def get_qlist(self):
        '''获取问题列表'''
        line = self.__fr.readline()
        while line:
            yield line.strip()
            line = self.__fr.readline()
        else:
            self.__fr.close()
    
    def to_jsonl(self, data: dict):
        '''将问答数据存入jsonlines'''
        json.dump(data, self.__fw, ensure_ascii=False)
        self.__fw.write('\n')
    
    def to_sql(self, data: dict):
        pass

    def close(self):
        if not self.__fr.closed:
            self.__fr.close()
        if not self.__fw.closed:
            self.__fw.close()


if __name__ == '__main__':

    # cookies = json.load(open('cookies.json', encoding='utf8'))['cookies']
    spider = FggSpider(ignore_certificate=True)
    iofiles = QAFiles('question/test.txt', 'test.jsonl')
    
    qlist = iofiles.get_qlist()
    spider.search_init()  # cookies=cookies)
    sys.exit(0)
    for q in qlist:
        print('start: %s' % q)
        try:
            result = {}
            result['question'] = q
            code = spider.send_question(q)
            answer = spider.get_answer(code)
            if answer:
                result.update(answer)
                iofiles.to_jsonl(result)
                print('%s done!' % q)
            else:
                print('%s has no answer' % q)
        except Exception as e:
            print('%s failed!' % q)
            continue
    else:
        iofiles.close()
        spider.close()
        print('done!')
